package cn.jetpiece.cloud.loadbalancer;

import cn.jetpiece.cloud.config.properties.JetpieceGatewayProperties;
import cn.jetpiece.cloud.core.consts.GatewayConstants;
import com.alibaba.cloud.commons.lang.StringUtils;
import com.alibaba.cloud.nacos.NacosDiscoveryProperties;
import com.alibaba.cloud.nacos.balancer.NacosBalancer;
import com.alibaba.cloud.nacos.loadbalancer.NacosLoadBalancer;
import com.alibaba.nacos.client.naming.utils.CollectionUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.*;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Mono;

import java.net.InetAddress;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * 自定义负载均衡、
 *
 * @author XuJZ
 */
@Slf4j
public class JetpieceNacosLoadBalancer extends NacosLoadBalancer implements ReactorServiceInstanceLoadBalancer {

    private final String serviceId;

    private final ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;

    private final NacosDiscoveryProperties nacosDiscoveryProperties;

    public JetpieceNacosLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider, String serviceId, NacosDiscoveryProperties nacosDiscoveryProperties) {
        super(serviceInstanceListSupplierProvider, serviceId, nacosDiscoveryProperties);
        this.serviceId = serviceId;
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
        this.nacosDiscoveryProperties = nacosDiscoveryProperties;
    }

    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        RequestDataContext dataContext = (RequestDataContext) request.getContext();
        RequestData requestData = dataContext.getClientRequest();
        ServiceInstanceListSupplier supplier = serviceInstanceListSupplierProvider.getIfAvailable(NoopServiceInstanceListSupplier::new);
        return supplier.get().next().mapNotNull(item -> this.getInstanceResponse(item, requestData));
    }

    private Response<ServiceInstance> getInstanceResponse(List<ServiceInstance> serviceInstances, RequestData requestData) {
        if (serviceInstances.isEmpty()) {
            log.warn("No servers available for service: " + this.serviceId);
            return new EmptyResponse();
        }
        try {
            String clusterName = this.nacosDiscoveryProperties.getClusterName();
            List<ServiceInstance> instancesToChoose = serviceInstances;
            if (StringUtils.isNotBlank(clusterName)) {
                List<ServiceInstance> sameClusterInstances = serviceInstances.stream().filter(serviceInstance -> {
                    String cluster = serviceInstance.getMetadata().get("nacos.cluster");
                    return StringUtils.equals(cluster, clusterName);
                }).collect(Collectors.toList());
                if (!CollectionUtils.isEmpty(sameClusterInstances)) {
                    instancesToChoose = sameClusterInstances;
                }
            } else {
                log.warn("A cross-cluster call occurs，name = {}, clusterName = {}, instance = {}", serviceId, clusterName, serviceInstances);
            }
            ServiceInstance instance;
            Map<String, String> routes = new HashMap<>();
            if (JetpieceGatewayProperties.getInstance().getLocalIpPreferenceRoute()) {
                String ip = getHostAddress();
                if (ip != null) {
                    routes.put(ip, ip);
                    instance = developerRouteInstance(instancesToChoose, routes, requestData);
                    if (instance != null) {
                        return new DefaultResponse(instance);
                    }
                }
            }
            routes = JetpieceGatewayProperties.getInstance().getDeveloperRoute().get(this.serviceId);
            if (routes != null && !routes.isEmpty()) {
                instance = developerRouteInstance(instancesToChoose, routes, requestData);
                if (instance != null) {
                    return new DefaultResponse(instance);
                }
            }
            instance = NacosBalancer.getHostByRandomWeight3(instancesToChoose);
            return new DefaultResponse(instance);
        } catch (Exception e) {
            log.warn("NacosLoadBalancer error " + e.getMessage(), e);
            return new EmptyResponse();
        }
    }

    /**
     * 获取开发路由实例
     * @param instancesToChoose 所有活动实例
     * @param routes 路由列表
     * @param requestData 请求数据
     * @return 命中实例
     */
    private ServiceInstance developerRouteInstance(List<ServiceInstance> instancesToChoose, Map<String, String> routes, RequestData requestData) {
        String ipAddress = requestData.getHeaders().getFirst(GatewayConstants.REAL_IP);
        if (!routes.containsKey(ipAddress)) {
            return null;
        }
        String target = routes.get(ipAddress);
        for (ServiceInstance serviceInstance : instancesToChoose) {
            if (serviceInstance.getHost().equals(target)) {
                return serviceInstance;
            }
        }
        return null;
    }

    /**
     * 缓存本机地址
     */
    private static String hostAddress;

    /**
     * 获取本机地址
     *
     * @return 本地地址
     */
    private String getHostAddress() {
        if (hostAddress == null) {
            try {
                hostAddress = InetAddress.getLocalHost().getHostAddress();
            } catch (Exception e) {
                log.error(e.getMessage(), e);
            }
        }
        return hostAddress;
    }
}
